#!/usr/bin/python3
import os
import cv2
import numpy as np
import pandas as pd

# face_attr_list_file = r'/rootfs/media/kasim/Data1/data/VideoCropFace/_quality1_attr.txt'
face_attr_list_file = r'/rootfs/media/kasim/Data1/data/VideoCropFace/face_image_filter_attr.txt'

USE_FILTER_DEVICE_IDS = True

gaze_angle_thr = 20.0

quality_thr = 0.20

min_pitch = -20
max_pitch = 20
min_roll = -30
max_roll = 30
min_yaw = -20
max_yaw = 20

min_angle_score = 0.1
max_angle_diff = 20

landmark_area_thr = 0.33
mouth_ratio_thr = 1.0

MASK_LEFT_EYEBROW_SCORE_THR     = 0.0
MASK_RIGHT_EYEBROW_SCORE_THR    = 0.0
MASK_LEFT_CHEEK_SCORE_THR       = 35.0
MASK_RIGHT_CHEEK_SCORE_THR      = 35.0
MASK_CHIN_CONTOUR_SCORE_THR     = 40.0
MASK_LEFT_EYE_SCORE_THR         = 20.0
MASK_RIGHT_EYE_SCORE_THR        = 20.0
MASK_MOUTH_SCORE_THR            = 35.0
MASK_NOSE_SCORE_THR             = 35.0

mask_score_thr = [
    MASK_LEFT_EYEBROW_SCORE_THR ,
    MASK_RIGHT_EYEBROW_SCORE_THR,
    MASK_LEFT_CHEEK_SCORE_THR   ,
    MASK_RIGHT_CHEEK_SCORE_THR  ,
    MASK_CHIN_CONTOUR_SCORE_THR ,
    MASK_LEFT_EYE_SCORE_THR     ,
    MASK_RIGHT_EYE_SCORE_THR    ,
    MASK_MOUTH_SCORE_THR        ,
    MASK_NOSE_SCORE_THR         ,
]

DEVICE_ID_SET = {
    '150100414a54443452067fa6d4556600',  # 104 529 19.66% 越秀星汇隽庭 7(幢/座) 1(单元) 1梯 A面
    '150100414a54443452067fa1d4096600',  #  76 391 19.44% 越秀星汇隽庭 4(幢/座) 1(单元) 2梯 A面
    '150100414a54443452064d27432e3600',  #  11  64 17.19% 大信时尚家园三期 12(幢/座) 1(单元) 1梯 A面
    '150100414a54443452064d2c43e13600',  #  33 198 16.67% 大信时尚家园二期 1(幢/座) 1(单元) 1梯 A面
    '150100414a5444345203bccb439fb500',  #  37 261 14.18% 越秀星汇隽庭 5(幢/座) 1(单元) 2梯 A面
    '150100414a54443452064d2c43ed3600',  #  70 496 14.11% 大信时尚家园三期 9(幢/座) 1(单元) 1梯 A面
    '150100414a54443452064d2742c73600',  #  52 370 14.05% 大信时尚家园二期 5(幢/座) 1(单元) 1梯 A面
    '150100414a54443452067fa6d45c6600',  #  51 399 12.78% 星光礼寓 7(幢/座) 1(单元) 1梯 A面
}


def main():
    # file_name;pitch;yaw;roll;pitch_bin;yaw_bin;roll_bin;gaze_angle;gaze;iqa;landmark_area;mouth_ratio;mask;landmark
    attr_cvs = pd.read_csv(face_attr_list_file, sep=';')

    def attr_statistics(attr_cvs, quality_thr):
        file_names = attr_cvs['file_name']
        qualitys = attr_cvs['iqa']
        bin_pitchs = attr_cvs['pitch_bin']
        bin_rolls = attr_cvs['roll_bin']
        bin_yaws = attr_cvs['yaw_bin']
        pitchs = attr_cvs['pitch']
        rolls = attr_cvs['roll']
        yaws = attr_cvs['yaw']
        gazes = attr_cvs['gaze']
        gaze_angles = attr_cvs['gaze_angle']
        landmark_areas = attr_cvs['landmark_area']
        mouth_ratios = attr_cvs['mouth_ratio']
        masks = attr_cvs['mask']
        # landmark = attr_cvs['landmark']

        # out_file_path = os.path.join(os.path.dirname(face_attr_list_file), 'image_list_good.attr.txt')           # 分类角度与回归角度差别较小的，且分类角度置信度较高的
        # out_file = open(out_file_path, 'w')
        # out_file.write('file_name;bin_pitch;bin_roll;bin_yaw;pitch;roll;yaw;quality')
        out_count = 0
        file_count = 0
        gaze_count = 0
        for file_name, quality_s, bin_pitch_s, bin_roll_s, bin_yaw_s, pitch_s, roll_s, yaw_s, gaze_s, gaze_angle_s, landmark_area_s, mouth_ratio_s, mask_s in zip(file_names, qualitys, bin_pitchs, bin_rolls, bin_yaws, pitchs, rolls, yaws, gazes, gaze_angles, landmark_areas, mouth_ratios, masks):
            file_count += 1
            # if file_count % 100 == 0:
            #     print('Process Filter File Count: {}'.format(file_count))

            gaze_angle = float(gaze_angle_s)
            if gaze_angle > gaze_angle_thr:
                continue

            gaze_count += 1
            # gaze = int(gaze_s)
            # if gaze == 0:
            #     continue

            quality = float(quality_s)
            if quality < quality_thr:
                continue

            pitch = float(pitch_s)
            roll = float(roll_s)
            yaw = float(yaw_s)

            if pitch < min_pitch or max_pitch < pitch or roll < min_roll or max_roll < roll or yaw < min_yaw or max_yaw < yaw:
                continue

            bin_pitch = bin_pitch_s.split(',')
            bin_roll = bin_roll_s.split(',')
            bin_yaw = bin_yaw_s.split(',')

            bin_pitch_score = float(bin_pitch[1])
            bin_pitch = float(bin_pitch[0])
            bin_roll_score = float(bin_roll[1])
            bin_roll = float(bin_roll[0])
            bin_yaw_score = float(bin_yaw[1])
            bin_yaw = float(bin_yaw[0])

            pitch_diff = abs(bin_pitch - pitch)
            roll_diff = abs(bin_roll - roll)
            yaw_diff = abs(bin_yaw - yaw)
            if max_angle_diff < pitch_diff or max_angle_diff < roll_diff or max_angle_diff < yaw_diff:
                 continue

            if bin_pitch_score < min_angle_score or bin_roll_score < min_angle_score or bin_yaw_score < min_angle_score:
                continue

            landmark_area = float(landmark_area_s)
            if landmark_area < landmark_area_thr:
                continue

            mouth_ratio = float(mouth_ratio_s)
            if mouth_ratio < mouth_ratio_thr:
                continue

            masks = mask_s.split(',')
            masks = list(map(float, masks))
            is_occ = False
            for i, mask in enumerate(masks):
                if mask < mask_score_thr[i]:
                    is_occ = True
                    break
            if is_occ:
                continue

            # out_info = '{};{};{};{};{};{};{};{}\n'.format(file_name, bin_pitch_s, bin_roll_s, bin_yaw_s, pitch_s, roll_s, yaw_s, quality_s)
            # out_file.write(out_info)
            out_count += 1
        print('Quality {:.02f}: Filter File Count: {}, Pass Count: {}, {:.02f}%, Gaze Count: {}, Gaze Count: {}, {:.02f}%'.format(quality_thr, file_count, out_count, 100*out_count/file_count, gaze_count, out_count, 100*out_count/gaze_count))
        # out_file.close()
        # os.system('chmod a+wr {}'.format(out_file_path))
        # print('out_count:', out_count)

    if USE_FILTER_DEVICE_IDS:
        file_names = attr_cvs['file_name']
        qualitys = attr_cvs['iqa']
        bin_pitchs = attr_cvs['pitch_bin']
        bin_rolls = attr_cvs['roll_bin']
        bin_yaws = attr_cvs['yaw_bin']
        pitchs = attr_cvs['pitch']
        rolls = attr_cvs['roll']
        yaws = attr_cvs['yaw']
        gazes = attr_cvs['gaze']
        gaze_angles = attr_cvs['gaze_angle']
        landmark_areas = attr_cvs['landmark_area']
        mouth_ratios = attr_cvs['mouth_ratio']
        masks = attr_cvs['mask']
        landmark = attr_cvs['landmark']

        _file_names = []
        _qualitys = []
        _bin_pitchs = []
        _bin_rolls = []
        _bin_yaws = []
        _pitchs = []
        _rolls = []
        _yaws = []
        _gazes = []
        _gaze_angles = []
        _landmark_areas = []
        _mouth_ratios = []
        _masks = []
        _landmark = []

        for i, file_name in enumerate(file_names):
            device_id = file_name.split('/')[-4]
            if device_id not in DEVICE_ID_SET:
                continue

            _file_names.append(file_names[i])
            _qualitys.append(qualitys[i])
            _bin_pitchs.append(bin_pitchs[i])
            _bin_rolls.append(bin_rolls[i])
            _bin_yaws.append(bin_yaws[i])
            _pitchs.append(pitchs[i])
            _rolls.append(rolls[i])
            _yaws.append(yaws[i])
            _gazes.append(gazes[i])
            _gaze_angles.append(gaze_angles[i])
            _landmark_areas.append(landmark_areas[i])
            _mouth_ratios.append(mouth_ratios[i])
            _masks.append(masks[i])
            _landmark.append(landmark[i])

        attr_cvs = {}
        attr_cvs['file_name'] = _file_names
        attr_cvs['iqa'] = _qualitys
        attr_cvs['pitch_bin'] = _bin_pitchs
        attr_cvs['roll_bin'] = _bin_rolls
        attr_cvs['yaw_bin'] = _bin_yaws
        attr_cvs['pitch'] = _pitchs
        attr_cvs['roll'] = _rolls
        attr_cvs['yaw'] = _yaws
        attr_cvs['gaze'] = _gazes
        attr_cvs['gaze_angle'] = _gaze_angles
        attr_cvs['landmark_area'] = _landmark_areas
        attr_cvs['mouth_ratio'] = _mouth_ratios
        attr_cvs['mask'] = _masks
        attr_cvs['landmark'] = _landmark

    # attr_statistics(attr_cvs, quality_thr)
    for _quality_thr in np.linspace(0.0, 0.5, 11):
        attr_statistics(attr_cvs, _quality_thr)


if __name__ == '__main__':
    main()
